#pragma once

#include <string>
#include <iostream>
#include <stdio.h>

#include "bytes_per_spmv.h"
#include "benchmark.h"

#include "../timer.h"
#include <cusp/multiply.h>
#include <cusp/graph/partition.h>
#include <cusp/multi_array1d.h>
#include <cusp/multi_matrix.h>
#include <cusp/detail/device/spmv/coo_flat_k.h>

int deviceCount;
cusp::array1d<int,cusp::host_memory> global_partition;

template <typename TestMatrix,  typename HostXArray, typename MultiArray1, typename MultiArray2>
void setupRowwiseMultiply( const TestMatrix& test_matrix, const HostXArray& hostXArray,
                           MultiArray1& multiXArray, MultiArray2& multiYArray, cusp::mult_method_type )
{
    typedef typename TestMatrix::value_type   ValueType;

    cusp::array1d<ValueType,cusp::host_memory> hostYArray(test_matrix.num_rows, 0);

    multiXArray = MultiArray1(test_matrix.partition, hostXArray);
    multiYArray = MultiArray2(test_matrix.partition, hostYArray);
}

template <typename TestMatrix,  typename HostXArray, typename MultiArray1, typename MultiArray2>
void setupRowwiseMultiply( const TestMatrix& test_matrix, const HostXArray& hostXArray,
                           MultiArray1& multiXArray, MultiArray2& multiYArray, cusp::multi_sync_mult )
{
    typedef typename TestMatrix::value_type   ValueType;

    cusp::array1d<ValueType,cusp::host_memory> hostYArray(test_matrix.num_rows, 0);

    multiXArray = MultiArray1(hostXArray);
    multiYArray = MultiArray2(test_matrix.partition, hostYArray);
}

template <typename TestMatrix, typename HostXArray, typename MultiArray1, typename MultiArray2>
void setupColwiseMultiply( const TestMatrix& test_matrix, const HostXArray& hostXArray,
                           MultiArray1& multiXArray, MultiArray2& multiYArray, cusp::mult_method_type )
{
    multiXArray = MultiArray1(test_matrix.partition, hostXArray);
    multiYArray = MultiArray2(test_matrix.num_rows, 0);
}

template <typename TestMatrix, typename HostXArray, typename MultiArray1, typename MultiArray2>
void setupMultiply( const TestMatrix& test_matrix, const HostXArray& hostXArray,
                    MultiArray1& multiXArray, MultiArray2& multiYArray, cusp::graph::colwise )
{
    /*cusp::csr_matrix<IndexType,ValueType,cusp::host_memory> csr(host_matrix);
    bool is_symmetric = cusp::graph::is_structurally_symmetric(csr);
    HostArray hostXArray(host_x);

    if( is_symmetric ) {
        cusp::permutation_matrix<int,cusp::host_memory> P;
        cusp::permutation_matrix<int,cusp::host_memory> Pt;
        cusp::graph::partitionToPermMatrix(*test_matrix.partition, P);
        cusp::transpose(P, Pt);
        cusp::multiply(Pt, host_x, hostXArray);
    }*/

    setupColwiseMultiply( test_matrix, hostXArray, multiXArray, multiYArray, typename TestMatrix::multi_mult_method() );
}

template <typename TestMatrix, typename HostXArray, typename MultiArray1, typename MultiArray2>
void setupMultiply( const TestMatrix& test_matrix, const HostXArray& hostXArray,
                    MultiArray1& multiXArray, MultiArray2& multiYArray, cusp::graph::rowwise )
{
    setupRowwiseMultiply( test_matrix, hostXArray, multiXArray, multiYArray, typename TestMatrix::multi_mult_method() );
}

template <typename TestMatrix, typename HostMatrix, typename HostArray, typename MultiArray1, typename MultiArray2>
void setupMultiply( const TestMatrix& test_matrix, const HostMatrix& host_matrix, const HostArray& hostXArray,
                    MultiArray1& multiXArray, MultiArray2& multiYArray )
{
    setupMultiply( test_matrix, hostXArray, multiXArray, multiYArray, typename TestMatrix::part_method() );
}

template <typename HostMatrix, typename TestMatrix>
float check_multi_spmv(HostMatrix& host_matrix, TestMatrix& test_matrix)
{
    typedef typename TestMatrix::index_type   IndexType; // ASSUME same as HostMatrix::index_type
    typedef typename TestMatrix::value_type   ValueType; // ASSUME same as HostMatrix::value_type
    typedef typename TestMatrix::memory_space MemorySpace;

    const IndexType M = host_matrix.num_rows;
    const IndexType N = host_matrix.num_cols;

    // create host input (x) and output (y) vectors
    cusp::array1d<ValueType,cusp::host_memory> host_x(N);
    cusp::array1d<ValueType,cusp::host_memory> host_y(M);
    //for(IndexType i = 0; i < N; i++) host_x[i] = (rand() % 21) - 10;
    for(IndexType i = 0; i < N; i++) host_x[i] = (int(i % 21) - 10);
    for(IndexType i = 0; i < M; i++) host_y[i] = 0;

    cusp::multi_array1d<ValueType, MemorySpace> test_x;
    cusp::multi_array1d<ValueType, MemorySpace> test_y;

    // create test input (x) and output (y) vectors
    setupMultiply( test_matrix, host_matrix, host_x, test_x, test_y );

    // compute SpMV on host and device
    cusp::multiply(host_matrix, host_x, host_y);
    cusp::multiply(test_matrix, test_x, test_y);

    // compare results
    cusp::array1d<ValueType,cusp::host_memory> test_y_copy(test_y);

    double error = l2_error(M, thrust::raw_pointer_cast(&test_y_copy[0]), thrust::raw_pointer_cast(&host_y[0]));

    return error;
}

template <typename HostMatrix, typename TestMatrix>
float time_multi_spmv(HostMatrix& host_matrix, TestMatrix& test_matrix, double seconds = 3.0, size_t min_iterations = 100, size_t max_iterations = 500)
{
    typedef typename TestMatrix::index_type   IndexType; // ASSUME same as HostMatrix::index_type
    typedef typename TestMatrix::value_type   ValueType; // ASSUME same as HostMatrix::value_type
    typedef typename TestMatrix::memory_space MemorySpace;
    typedef typename TestMatrix::RemoteContainer::memory_space RemoteSpace;

    // create test input (x) and output (y) vectors
    size_t N = test_matrix.num_cols;

    cusp::array1d<IndexType,cusp::host_memory> x(N);

    cusp::multi_array1d<ValueType, MemorySpace> test_x;
    cusp::multi_array1d<ValueType, MemorySpace> test_y;

    // create test input (x) and output (y) vectors
    setupMultiply( test_matrix, host_matrix, x, test_x, test_y );

    // burn one extra call for initialization of threads
    cusp::multiply(test_matrix, test_x, test_y);

    // warmup
    timer time_one_iteration;
    cusp::multiply(test_matrix, test_x, test_y);
    cudaThreadSynchronize();
    double estimated_time = time_one_iteration.seconds_elapsed();

    // determine # of seconds dynamically
    size_t num_iterations;
    if (estimated_time == 0)
        num_iterations = max_iterations;
    else
        num_iterations = std::min(max_iterations, std::max(min_iterations, (size_t) (seconds / estimated_time)) );

    // time several SpMV iterations
    timer t;
    for(size_t i = 0; i < num_iterations; i++)
        cusp::multiply(test_matrix, test_x, test_y);

    cudaThreadSynchronize();

    float sec_per_iteration = t.seconds_elapsed() / num_iterations;

    return sec_per_iteration;
}

template <typename HostMatrix, typename TestMatrixOnHost, typename TestMatrixOnDevice>
void test_multi_spmv(std::string   kernel_name,
                     HostMatrix&         host_matrix,
                     TestMatrixOnHost&   test_matrix_on_host,
                     TestMatrixOnDevice& test_matrix_on_device)
{
    float error = check_multi_spmv(host_matrix, test_matrix_on_device);
    float time  = time_multi_spmv(host_matrix, test_matrix_on_device);
    float gbyte = bytes_per_spmv(test_matrix_on_host);

    float GFLOPs = (time == 0) ? 0 : (2 * host_matrix.num_entries / time) / 1e9;
    float GBYTEs = (time == 0) ? 0 : (gbyte / time)                       / 1e9;

    printf("\t%-20s: %8.4f ms ( %5.2f GFLOP/s %5.1f GB/s) [L2 error %f]\n", kernel_name.c_str(), 1e3 * time, GFLOPs, GBYTEs, error);

    //record results to file
    FILE * fid = fopen(BENCHMARK_OUTPUT_FILE_NAME, "a");
    fprintf(fid, "kernel=%s gflops=%f gbytes=%f msec=%f num_devices=%d\n", kernel_name.c_str(), GFLOPs, GBYTEs, 1e3 * time, deviceCount);
    fclose(fid);
}


/////////////////////////////////////////////////////
// These methods test specific formats and kernels //
/////////////////////////////////////////////////////
template <typename Partitioner, typename PartMethod, typename HostMatrix>
void setup_multi_vars( HostMatrix& test_matrix_on_host, size_t testDeviceNum)
{
    deviceCount = testDeviceNum;

    timer t;
    cusp::graph::partition<Partitioner,PartMethod>(test_matrix_on_host, deviceCount, global_partition);
    printf("Partition time : %f\n", t.milliseconds_elapsed());
}

template <typename Partitioner, typename PartMethod, typename HostMatrix>
void test_multi_sync(HostMatrix& host_matrix)
{
    typedef typename HostMatrix::index_type IndexType;
    typedef typename HostMatrix::value_type ValueType;
    typedef typename cusp::multi_array1d<ValueType,cusp::device_memory> MultiArray;

    // convert HostMatrix to TestMatrix on host
    cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> test_matrix_on_host(host_matrix);

    // transfer TestMatrix to device
    typedef typename cusp::multi_csr_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_sync_mult, Partitioner, PartMethod> CsrDeviceMatrix;
    typedef typename cusp::multi_coo_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_sync_mult, Partitioner, PartMethod> CooDeviceMatrix;
    typedef typename cusp::multi_ell_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_sync_mult, Partitioner, PartMethod> EllDeviceMatrix;
    typedef typename cusp::multi_hyb_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_sync_mult, Partitioner, PartMethod> HybDeviceMatrix;

    CooDeviceMatrix test_coo_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_coo_sync",     host_matrix, test_matrix_on_host, test_coo_matrix_on_device);
    CsrDeviceMatrix test_csr_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_csr_sync",     host_matrix, test_matrix_on_host, test_csr_matrix_on_device);

    try
    {
        EllDeviceMatrix test_ell_matrix_on_device(test_matrix_on_host, global_partition);
        test_multi_spmv("multi_ell_sync",     host_matrix, test_matrix_on_host, test_ell_matrix_on_device);
    }
    catch (cusp::format_conversion_exception)
    {
        std::cout << "\tRefusing to convert to multi ELL format" << std::endl;
    }

    HybDeviceMatrix test_hyb_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_hyb_sync",     host_matrix, test_matrix_on_host, test_hyb_matrix_on_device);
}

template <typename Partitioner, typename PartMethod, typename HostMatrix>
void test_multi_kernel(HostMatrix& host_matrix)
{
    typedef typename HostMatrix::index_type IndexType;
    typedef typename HostMatrix::value_type ValueType;
    typedef typename cusp::multi_array1d<ValueType,cusp::device_memory> MultiArray;

    // convert HostMatrix to TestMatrix on host
    cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> test_matrix_on_host(host_matrix);

    // transfer TestMatrix to device
    typedef typename cusp::multi_csr_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_kernel_mult, Partitioner, PartMethod> CsrDeviceMatrix;
    typedef typename cusp::multi_coo_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_kernel_mult, Partitioner, PartMethod> CooDeviceMatrix;
    typedef typename cusp::multi_ell_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_kernel_mult, Partitioner, PartMethod> EllDeviceMatrix;
    typedef typename cusp::multi_hyb_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_kernel_mult, Partitioner, PartMethod> HybDeviceMatrix;

    CooDeviceMatrix test_coo_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_coo_kernel",     host_matrix, test_matrix_on_host, test_coo_matrix_on_device);
    //CsrDeviceMatrix test_csr_matrix_on_device(test_matrix_on_host, global_partition);
    //test_multi_spmv("multi_csr_kernel",     host_matrix, test_matrix_on_host, test_csr_matrix_on_device);

    try
    {
        EllDeviceMatrix test_ell_matrix_on_device(test_matrix_on_host, global_partition);
        test_multi_spmv("multi_ell_kernel",     host_matrix, test_matrix_on_host, test_ell_matrix_on_device);
    }
    catch (cusp::format_conversion_exception)
    {
        std::cout << "\tRefusing to convert to multi ELL format" << std::endl;
    }

    HybDeviceMatrix test_hyb_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_hyb_kernel",     host_matrix, test_matrix_on_host, test_hyb_matrix_on_device);
}

template <typename Partitioner, typename PartMethod, typename HostMatrix>
void test_multi_host(HostMatrix& host_matrix)
{
    typedef typename HostMatrix::index_type IndexType;
    typedef typename HostMatrix::value_type ValueType;
    typedef typename cusp::multi_array1d<ValueType,cusp::device_memory> MultiArray;

    // convert HostMatrix to TestMatrix on host
    cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> test_matrix_on_host(host_matrix);

    // transfer TestMatrix to device
    typedef typename cusp::multi_csr_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_host_mult, Partitioner, PartMethod> CsrDeviceMatrix;
    typedef typename cusp::multi_coo_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_host_mult, Partitioner, PartMethod> CooDeviceMatrix;
    typedef typename cusp::multi_ell_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_host_mult, Partitioner, PartMethod> EllDeviceMatrix;
    typedef typename cusp::multi_hyb_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_host_mult, Partitioner, PartMethod> HybDeviceMatrix;

    CooDeviceMatrix test_coo_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_coo_host",     host_matrix, test_matrix_on_host, test_coo_matrix_on_device);
    //CsrDeviceMatrix test_csr_matrix_on_device(test_matrix_on_host, global_partition);
    //test_multi_spmv("multi_csr_host",     host_matrix, test_matrix_on_host, test_csr_matrix_on_device);

    try
    {
        EllDeviceMatrix test_ell_matrix_on_device(test_matrix_on_host, global_partition);
        test_multi_spmv("multi_ell_host",     host_matrix, test_matrix_on_host, test_ell_matrix_on_device);
    }
    catch (cusp::format_conversion_exception)
    {
        std::cout << "\tRefusing to convert to multi ELL format" << std::endl;
    }

    HybDeviceMatrix test_hyb_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_hyb_host",     host_matrix, test_matrix_on_host, test_hyb_matrix_on_device);
}

template <typename Partitioner, typename PartMethod, typename HostMatrix>
void test_multi_p2p(HostMatrix& host_matrix)
{
    typedef typename HostMatrix::index_type IndexType;
    typedef typename HostMatrix::value_type ValueType;
    typedef typename cusp::multi_array1d<ValueType,cusp::device_memory> MultiArray;

    // convert HostMatrix to TestMatrix on host
    cusp::coo_matrix<IndexType, ValueType, cusp::host_memory> test_matrix_on_host(host_matrix);

    // transfer TestMatrix to device
    typedef typename cusp::multi_csr_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_p2p_copy_mult, Partitioner, PartMethod> CsrDeviceMatrix;
    typedef typename cusp::multi_coo_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_p2p_copy_mult, Partitioner, PartMethod> CooDeviceMatrix;
    typedef typename cusp::multi_ell_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_p2p_copy_mult, Partitioner, PartMethod> EllDeviceMatrix;
    typedef typename cusp::multi_hyb_matrix<IndexType, ValueType, cusp::device_memory, cusp::multi_p2p_copy_mult, Partitioner, PartMethod> HybDeviceMatrix;

    CooDeviceMatrix test_coo_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_coo_p2p_copy",     host_matrix, test_matrix_on_host, test_coo_matrix_on_device);
    //CsrDeviceMatrix test_csr_matrix_on_device(test_matrix_on_host, global_partition);
    //test_multi_spmv("multi_csr_p2p_copy",     host_matrix, test_matrix_on_host, test_csr_matrix_on_device);

    try
    {
        EllDeviceMatrix test_ell_matrix_on_device(test_matrix_on_host, global_partition);
        test_multi_spmv("multi_ell_p2p_copy",     host_matrix, test_matrix_on_host, test_ell_matrix_on_device);
    }
    catch (cusp::format_conversion_exception)
    {
        std::cout << "\tRefusing to convert to multi ELL format" << std::endl;
    }

    HybDeviceMatrix test_hyb_matrix_on_device(test_matrix_on_host, global_partition);
    test_multi_spmv("multi_hyb_p2p_copy",     host_matrix, test_matrix_on_host, test_hyb_matrix_on_device);
}

